{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "import open3d as o3d\n",
    "import numpy as np\n",
    "import copy\n",
    "import time\n",
    "import os\n",
    "import sys\n",
    "\n",
    "# monkey patches visualization and provides helpers to load geometries\n",
    "sys.path.append('..')\n",
    "import open3d_tutorial as o3dtut\n",
    "# change to True if you want to interact with the visualization windows\n",
    "o3dtut.interactive = not \"CI\" in os.environ"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Global registration\n",
    "Both [ICP registration](../pipelines/icp_registration.ipynb) and [Colored point cloud registration](colored_pointcloud_registration.ipynb) are known as local registration methods because they rely on a rough alignment as initialization. This tutorial shows another class of registration methods, known as **global** registration. This family of algorithms do not require an alignment for initialization. They usually produce less tight alignment results and are used as initialization of the local methods."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization\n",
    "This helper function visualizes the transformed source point cloud together with the target point cloud:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_registration_result(source, target, transformation):\n",
    "    source_temp = copy.deepcopy(source)\n",
    "    target_temp = copy.deepcopy(target)\n",
    "    source_temp.paint_uniform_color([1, 0.706, 0])\n",
    "    target_temp.paint_uniform_color([0, 0.651, 0.929])\n",
    "    source_temp.transform(transformation)\n",
    "    o3d.visualization.draw_geometries([source_temp, target_temp],\n",
    "                                      zoom=0.4559,\n",
    "                                      front=[0.6452, -0.3036, -0.7011],\n",
    "                                      lookat=[1.9892, 2.0208, 1.8945],\n",
    "                                      up=[-0.2779, -0.9482, 0.1556])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract geometric feature\n",
    "We downsample the point cloud, estimate normals, then compute a FPFH feature for each point. The FPFH feature is a 33-dimensional vector that describes the local geometric property of a point. A nearest neighbor query in the 33-dimensinal space can return points with similar local geometric structures. See [\\[Rasu2009\\]](../reference.html#rasu2009) for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_point_cloud(pcd, voxel_size):\n",
    "    print(\":: Downsample with a voxel size %.3f.\" % voxel_size)\n",
    "    pcd_down = pcd.voxel_down_sample(voxel_size)\n",
    "\n",
    "    radius_normal = voxel_size * 2\n",
    "    print(\":: Estimate normal with search radius %.3f.\" % radius_normal)\n",
    "    pcd_down.estimate_normals(\n",
    "        o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=30))\n",
    "\n",
    "    radius_feature = voxel_size * 5\n",
    "    print(\":: Compute FPFH feature with search radius %.3f.\" % radius_feature)\n",
    "    pcd_fpfh = o3d.pipelines.registration.compute_fpfh_feature(\n",
    "        pcd_down,\n",
    "        o3d.geometry.KDTreeSearchParamHybrid(radius=radius_feature, max_nn=100))\n",
    "    return pcd_down, pcd_fpfh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input\n",
    "This code below reads a source point cloud and a target point cloud from two files. They are misaligned with an identity matrix as transformation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_dataset(voxel_size):\n",
    "    print(\":: Load two point clouds and disturb initial pose.\")\n",
    "\n",
    "    demo_icp_pcds = o3d.data.DemoICPPointClouds()\n",
    "    source = o3d.io.read_point_cloud(demo_icp_pcds.paths[0])\n",
    "    target = o3d.io.read_point_cloud(demo_icp_pcds.paths[1])\n",
    "    trans_init = np.asarray([[0.0, 0.0, 1.0, 0.0], [1.0, 0.0, 0.0, 0.0],\n",
    "                             [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]])\n",
    "    source.transform(trans_init)\n",
    "    draw_registration_result(source, target, np.identity(4))\n",
    "\n",
    "    source_down, source_fpfh = preprocess_point_cloud(source, voxel_size)\n",
    "    target_down, target_fpfh = preprocess_point_cloud(target, voxel_size)\n",
    "    return source, target, source_down, target_down, source_fpfh, target_fpfh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "voxel_size = 0.05  # means 5cm for this dataset\n",
    "source, target, source_down, target_down, source_fpfh, target_fpfh = prepare_dataset(\n",
    "    voxel_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RANSAC\n",
    "We use RANSAC for global registration. In each RANSAC iteration, `ransac_n` random points are picked from the source point cloud. Their corresponding points in the target point cloud are detected by querying the nearest neighbor in the 33-dimensional FPFH feature space. A pruning step takes fast pruning algorithms to quickly reject false matches early.\n",
    "\n",
    "Open3D provides the following pruning algorithms:\n",
    "\n",
    "- `CorrespondenceCheckerBasedOnDistance` checks if aligned point clouds are close (less than the specified threshold).\n",
    "- `CorrespondenceCheckerBasedOnEdgeLength` checks if the lengths of any two arbitrary edges (line formed by two vertices) individually drawn from source and target correspondences are similar. This tutorial checks that $||edge_{source}|| > 0.9 \\cdot ||edge_{target}||$ and $||edge_{target}|| > 0.9 \\cdot ||edge_{source}||$ are true.\n",
    "- `CorrespondenceCheckerBasedOnNormal` considers vertex normal affinity of any correspondences. It computes the dot product of two normal vectors. It takes a radian value for the threshold.\n",
    "\n",
    "Only matches that pass the pruning step are used to compute a transformation, which is validated on the entire point cloud. The core function is `registration_ransac_based_on_feature_matching`. The most important hyperparameter of this function is `RANSACConvergenceCriteria`. It defines the maximum number of RANSAC iterations and the confidence probability. The larger these two numbers are, the more accurate the result is, but also the more time the algorithm takes.\n",
    "\n",
    "We set the RANSAC parameters based on the empirical value provided by [\\[Choi2015]\\](../reference.html#choi2015). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def execute_global_registration(source_down, target_down, source_fpfh,\n",
    "                                target_fpfh, voxel_size):\n",
    "    distance_threshold = voxel_size * 1.5\n",
    "    print(\":: RANSAC registration on downsampled point clouds.\")\n",
    "    print(\"   Since the downsampling voxel size is %.3f,\" % voxel_size)\n",
    "    print(\"   we use a liberal distance threshold %.3f.\" % distance_threshold)\n",
    "    result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(\n",
    "        source_down, target_down, source_fpfh, target_fpfh, True,\n",
    "        distance_threshold,\n",
    "        o3d.pipelines.registration.TransformationEstimationPointToPoint(False),\n",
    "        3, [\n",
    "            o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(\n",
    "                0.9),\n",
    "            o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(\n",
    "                distance_threshold)\n",
    "        ], o3d.pipelines.registration.RANSACConvergenceCriteria(100000, 0.999))\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_ransac = execute_global_registration(source_down, target_down,\n",
    "                                            source_fpfh, target_fpfh,\n",
    "                                            voxel_size)\n",
    "print(result_ransac)\n",
    "draw_registration_result(source_down, target_down, result_ransac.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "    \n",
    "**Note:**\n",
    "\n",
    "Open3D provides a faster implementation for global registration. Please refer to [Fast global registration](#fast-global-registration).\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Local refinement\n",
    "For performance reason, the global registration is only performed on a heavily down-sampled point cloud. The result is also not tight. We use [Point-to-plane ICP](../icp_registration.ipynb#point-to-plane-ICP) to further refine the alignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def refine_registration(source, target, source_fpfh, target_fpfh, voxel_size):\n",
    "    distance_threshold = voxel_size * 0.4\n",
    "    print(\":: Point-to-plane ICP registration is applied on original point\")\n",
    "    print(\"   clouds to refine the alignment. This time we use a strict\")\n",
    "    print(\"   distance threshold %.3f.\" % distance_threshold)\n",
    "    result = o3d.pipelines.registration.registration_icp(\n",
    "        source, target, distance_threshold, result_ransac.transformation,\n",
    "        o3d.pipelines.registration.TransformationEstimationPointToPlane())\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_icp = refine_registration(source, target, source_fpfh, target_fpfh,\n",
    "                                 voxel_size)\n",
    "print(result_icp)\n",
    "draw_registration_result(source, target, result_icp.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fast global registration\n",
    "The RANSAC based global registration solution may take a long time due to countless model proposals and evaluations. [\\[Zhou2016\\]](../reference.html#zhou2016) introduced a faster approach that quickly optimizes line process weights of few correspondences. As there is no model proposal and evaluation involved for each iteration, the approach proposed in [\\[Zhou2016\\]](../reference.html#zhou2016) can save a lot of computational time.\n",
    "\n",
    "This tutorial compares the running time of the RANSAC based global registration to the implementation of [\\[Zhou2016\\]](../reference.html#zhou2016)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input\n",
    "We use the same input as in the global registration example above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "voxel_size = 0.05  # means 5cm for the dataset\n",
    "source, target, source_down, target_down, source_fpfh, target_fpfh = \\\n",
    "        prepare_dataset(voxel_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Baseline\n",
    "In the code below we time the global registration approach."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "result_ransac = execute_global_registration(source_down, target_down,\n",
    "                                            source_fpfh, target_fpfh,\n",
    "                                            voxel_size)\n",
    "print(\"Global registration took %.3f sec.\\n\" % (time.time() - start))\n",
    "print(result_ransac)\n",
    "draw_registration_result(source_down, target_down, result_ransac.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fast global registration\n",
    "With the same input used for a baseline, the code below calls the implementation of [\\[Zhou2016\\]](../reference.html#zhou2016)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def execute_fast_global_registration(source_down, target_down, source_fpfh,\n",
    "                                     target_fpfh, voxel_size):\n",
    "    distance_threshold = voxel_size * 0.5\n",
    "    print(\":: Apply fast global registration with distance threshold %.3f\" \\\n",
    "            % distance_threshold)\n",
    "    result = o3d.pipelines.registration.registration_fgr_based_on_feature_matching(\n",
    "        source_down, target_down, source_fpfh, target_fpfh,\n",
    "        o3d.pipelines.registration.FastGlobalRegistrationOption(\n",
    "            maximum_correspondence_distance=distance_threshold))\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "result_fast = execute_fast_global_registration(source_down, target_down,\n",
    "                                               source_fpfh, target_fpfh,\n",
    "                                               voxel_size)\n",
    "print(\"Fast global registration took %.3f sec.\\n\" % (time.time() - start))\n",
    "print(result_fast)\n",
    "draw_registration_result(source_down, target_down, result_fast.transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With proper configuration, the accuracy of fast global registration is even comparable to ICP. Please refer to [\\[Zhou2016\\]](../reference.html#zhou2016) for more experimental results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to FPFH-feature-based FGR, global registration can be performed with correspondence-based FGR via `registration_fgr_based_on_correspondence`. This method is useful if your correspondence frontend is different than FPFH, but you would still like to use FGR given a set of putative correspondences. It can be called with\n",
    "\n",
    "```python\n",
    "o3d.pipelines.registration.registration_fgr_based_on_correspondence(\n",
    "        source_down, target_down, correspondence_set,\n",
    "        o3d.pipelines.registration.FastGlobalRegistrationOption())\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Edit Metadata",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
